import mysql.connector
import json
import os
from contextlib import contextmanager
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

# 加载环境变量
load_dotenv()


class DbManager:
    def __init__(self):
        self.connection_pool = mysql.connector.pooling.MySQLConnectionPool(
            pool_name="db_pool",
            pool_size=5,
            pool_reset_session=True,
            host=os.getenv("DB_HOST"),  # 数据库服务器地址
            user=os.getenv("DB_USER"),  # 数据库用户名
            password=os.getenv("DB_PASSWD"),  # 数据库密码
            database=os.getenv("DB_NAME"),  # 数据库名
        )

    @contextmanager
    def get_cursor(self):
        with self.connection_pool.get_connection() as connection:
            cursor = None
            try:
                cursor = connection.cursor()
                yield cursor
                connection.commit()
            except Exception as e:
                connection.rollback()
                raise e
            finally:
                if cursor:
                    cursor.close()

    def execute_sql(self, sql: str) -> str:
        print("execute_sql开始执行%s" % sql)
        with self.get_cursor() as cursor:
            cursor.execute(sql)
            if cursor.description is not None:
                rows = cursor.fetchall()
                result = {
                    "columns": [desc[0] for desc in cursor.description],
                    "rows": rows,
                }
                return json.dumps(result, default=str)
            else:
                return f"row affected:{cursor.rowcount}"


dbManager = DbManager()
# Initialize FastMCP server
mcp = FastMCP(name="finance_manager", port=9999, host="0.0.0.0")


@mcp.tool()
def get_table_definition(table: str) -> str:
    print("get_table_definition获取表%s的结构定义" % table)
    """get table definition"""
    return dbManager.execute_sql(f"show create table {table}")


@mcp.tool()
def execute_sql(sql: str) -> str:
    """execute sql"""
    return dbManager.execute_sql(sql)


if __name__ == "__main__":
    # mcp.run(transport="sse")
    mcp.run(transport="sse")
